import cv2
import numpy as np
import subprocess               # 用于启动和管理FFmpeg子进程
import threading
from queue import Queue, Empty  # 线程安全的队列，用于帧传递
import time
from typing import Optional, Tuple
import os


class RTSPPublisher:
    def __init__(self, config_loader):
        self.config = config_loader.get_app_config()        # 1. 读取推流相关配置
        # 使用不带认证的URL进行推流
        self.output_url = self.config['rtsp']['output_url']  
        # 保存带认证的URL供显示使用
        self.view_url = self.config['rtsp'].get('view_url', self.output_url)
        self.video_config = self.config['video']            # 4. 获取视频参数（分辨率、帧率、编码等）

        self.frame_queue = Queue(maxsize=30)                # 5. 初始化帧队列，最大缓存30帧，防止内存溢出
        self.is_running = False                             # 推流服务运行状态
        self.publish_thread = None                          # 推流线程对象
        self.ffmpeg_process = None                          # 当前运行的FFmpeg进程

        # 推流参数
        self.width = self.video_config['frame_width']       # 视频宽度
        self.height = self.video_config['frame_height']     # 视频高度
        self.fps = self.video_config['fps']                 # 帧率
        self.codec = self.video_config['codec']             # 编码器

    def start(self) -> bool:
        """启动RTSP推流服务"""
        try:
            # 构建FFmpeg命令 - 注意这里使用不带认证的URL
            ffmpeg_cmd = [
                'ffmpeg',
                '-y',                       # 覆盖输出文件（如果有）
                '-f', 'rawvideo',           # 输入格式为原始视频流
                '-vcodec', 'rawvideo',      # 输入流的编码格式
                '-pix_fmt', 'bgr24',        # 输入像素格式（OpenCV默认BGR）
                '-s', f'{self.width}x{self.height}',    # 输入分辨率
                '-r', str(self.fps),        # 输入帧率
                '-i', '-',                  # 从stdin读取，即从标准输入读取视频流
                '-c:v', 'libx264',          # 使用H.264编码
                '-preset', 'ultrafast',     # 最快的编码预设，减少延迟，提高编码效率
                '-tune', 'zerolatency',     # 零延迟调优
                '-pix_fmt', 'yuv420p',      # 输出像素格式，兼容性好
                '-g', '15',      # 输出像素格式，兼容性好
                '-f', 'rtsp',               # 输出格式为RTSP
                '-rtsp_transport', 'tcp',   # 使用TCP传输
                self.output_url             # 推流到MediaMTX（不需要认证）
            ]
            print("FFmpeg命令:", ffmpeg_cmd)
            # 启动FFmpeg进程，stdin用于写入帧数据，stderr用于捕获错误信息
            self.ffmpeg_process = subprocess.Popen(
                ffmpeg_cmd,
                stdin=subprocess.PIPE,      # 主线程通过stdin写入帧
                stderr=subprocess.PIPE,     # 捕获FFmpeg的错误输出
                bufsize=10 ** 8             # 设置较大的缓冲区，防止阻塞
            )

            self.is_running = True          # 标记推流服务为运行状态

             # 3. 启动推流线程，负责不断从队列取帧并写入FFmpeg
            self.publish_thread = threading.Thread(target=self._publish_frames)
            self.publish_thread.daemon = True       # 设置为守护线程，主程序退出时自动退出
            self.publish_thread.start()             # 启动推流线程

            # 显示带认证的观看URL
            print(f"RTSP推流服务已启动")
            print(f"推流地址: {self.output_url}")
            print(f"观看地址: {self.view_url}")
            return True

        except Exception as e:
            print(f"启动RTSP推流失败: {e}")
            return False

    def get_view_url(self) -> str:
        """获取带认证的观看URL（供Web界面显示用）"""
        return self.view_url

    def _publish_frames(self):
        """推送帧的线程函数，不断从队列取帧并写入FFmpeg进程"""
        while self.is_running:   # 只要推流服务处于运行状态
            try:
                # 从帧队列中取出一帧，最多等待0.1秒（100ms），防止线程阻塞
                frame = self.frame_queue.get(timeout=0.1)

                # 检查帧有效、FFmpeg进程存在且未退出
                if frame is not None and self.ffmpeg_process and self.ffmpeg_process.poll() is None:
                    # 确保帧大小正确，如果帧大小与配置不一致，则调整
                    if frame.shape[:2] != (self.height, self.width):
                        frame = cv2.resize(frame, (self.width, self.height))    # 调整帧大小

                    # 写入FFmpeg进程的标准输入
                    try:
                        self.ffmpeg_process.stdin.write(frame.tobytes())    # 将帧数据写入FFmpeg进程的标准输入
                        self.ffmpeg_process.stdin.flush()                   # 刷新缓冲区，确保数据及时发送
                    except BrokenPipeError:
                        print("FFmpeg进程已断开，尝试重启...")
                        self._restart_ffmpeg()

            except Empty:
                continue                    # 如果队列为空，则继续循环
            except Exception as e:
                print(f"推流错误: {e}")       # 打印错误信息
                time.sleep(0.1)             # 等待0.1秒后继续循环，避免死循环占用cpu

    def _restart_ffmpeg(self):
        """重启FFmpeg进程（如推流中断时自动恢复）"""
        try:
            if self.ffmpeg_process:
                self.ffmpeg_process.terminate()         # 终止FFmpeg进程
                self.ffmpeg_process.wait(timeout=5)     # 等待5秒，确保进程完全退出
            self.start()            # 重新启动推流服务
        except Exception as e:
            print(f"重启FFmpeg失败: {e}")

    def push_frame(self, frame: np.ndarray):
        """推送单帧"""
        try:
            # 如果队列中帧数超过10，清空旧帧，保证低延迟（只保留最新帧）
            while self.frame_queue.qsize() > 10:
                try:
                    self.frame_queue.get_nowait()
                except:
                    break

            self.frame_queue.put(frame)     # 将帧添加到队列
        except:
            pass        # 忽略所有异常，保证主流程不被阻塞

    def stop(self):
        """停止推流服务"""
        self.is_running = False                     # 标记推流服务为停止状态

        if self.publish_thread:
            self.publish_thread.join(timeout=2)     # 等待推流线程退出，最多2秒

        if self.ffmpeg_process:
            self.ffmpeg_process.terminate()         # 终止FFmpeg进程
            try:
                self.ffmpeg_process.wait(timeout=5)  # 等待5秒，确保进程完全退出
            except:
                self.ffmpeg_process.kill()           # 如果超时，则强制终止进程

        print("RTSP推流服务已停止")     
